import os, cv2
import numpy as np
from math import ceil
from dataset.data_aug import *
import time
import torch
from models.MobileNet import MobileNet,mobilenet_v2
import matplotlib.pyplot as plt
from torchvision import transforms
from PIL import Image

def class2_init(resume):
    test_transforms= Compose([
                                # ExpandBorder(size=(48, 26), resize=True),
                              # transforms.ToTensor(),
                              Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    # model = MobileNet(2, alpha=0.5)
    model=mobilenet_v2(2)
    # model.fc = torch.nn.Linear(100352, 2)
    # model.fc = torch.nn.Linear(model.fc.in_features, 2)
    model.load_state_dict(torch.load(resume), strict=True)
    model = model.cuda()
    model.eval()
    return model, test_transforms


def class2_forward(testimage_list, class_model_list):

    # 1 涂卡
    # 0 没涂卡
    class_model, class_transforms = class_model_list
    class_model.eval()
    input = []
    resize=224
    for i in range(len(testimage_list)):
        # temp_small = cv2.cvtColor(cv2.resize(testimage_list[i], (48, 26)), cv2.COLOR_GRAY2RGB)
        # min_scale = min(resize/ testimage_list[i].shape[0], resize/ testimage_list[i].shape[1])
        # temp_small = cv2.cvtColor(cv2.resize(testimage_list[i], (0,0),fx=min_scale,fy=min_scale), cv2.COLOR_BGR2RGB)
        # white_img = np.zeros((resize, resize, 3), dtype=np.uint8) + 255
        # real_img = white_img.copy()
        # real_img[:temp_small.shape[0], :temp_small.shape[1], :] = temp_small
        # temp_small_ = real_img

        temp_small = cv2.cvtColor(cv2.resize(testimage_list[i], (224,224)), cv2.COLOR_BGR2RGB)
        one_input = class_transforms(temp_small)

        input.append(one_input)
    input = torch.from_numpy(np.asarray(input)).float()

    input = input.to(torch.device("cuda"))
    output = class_model(input)
    # props, preds = torch.max(torch.sigmoid(output), 1)
    # test_preds = preds.data.cpu().numpy()
    # test_scores = output.data.cpu().numpy()
    return output


if __name__ == '__main__':

    #########   set the GPU   ###########
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    #########  the input path   ##########
    import glob
    import tqdm,shutil



    # bar = tqdm.tqdm()
    resume = '../save_models/20220506_output/weights-5-0-1.0.pth'
    model, test_transforms = class2_init(resume)

    right = 0
    error = 0
    # img_list = glob.glob('/home/hegang/datas2/hegang/datas/meter/class_data/single_double_bj/1/*.jpg')
    img_list = glob.glob('/home/hegang/datas2/hegang/datas/meter/class_data/4_meters/1/*.jpg')
    # img_list=glob.glob(r"/data2/hjm/1/*.jpg")[0:1000]
    pbar=tqdm.tqdm(img_list)
    # print(len(img_list))
    # exit()
    # save_path=r"/data2/enducation/datas/answer_card/class_2/20211122/data/"
    # save_path0 = os.path.join(save_path, "0")
    # save_path1 = os.path.join(save_path, "1")
    # for save_path in [save_path0, save_path1]:
    #     if not os.path.exists(save_path):
    #         os.makedirs(save_path)

    for num, item in enumerate(pbar):
        print(item)
        # item = '/media/heils_lhl/data/youirobotData/light/2/1577264278.065616_ .jpg'
        img_ori = cv2.imread(item)
        # img = cv2.cvtColor(img_ori,cv2.COLOR_BGR2RGB)
        img = img_ori
        testimage = [img]
        result = class2_forward(testimage,[model,test_transforms])
        # print(result)
        # gt = int(item.split('/')[-2])
        pre,index=torch.max(result,1)

        print(index)
        import time
        # time.sleep(1)
        # if index==torch.tensor(0):
        #     shutil.copy(item,os.path.join(save_path0, os.path.split(item)[1]))
        # else:
        #     shutil.copy(item, os.path.join(save_path1,  os.path.split(item)[1]))
        # print(item)


    #     if gt != prd:
    #         error += 1
    #         print('result: %d, gt: %d, item: %s' % (result, gt, item))
    #
    #         targretfolder = '/data1/sheng/temp/1211/2_error _bozhan/'
    #         if not os.path.exists(targretfolder):
    #             os.makedirs(targretfolder)
    #         # filename = item.split('/')[-1]
    #         filename = "gt_{}_pred_{}_{}.jpg".format(gt, prd, num)
    #         cv2.imwrite(targretfolder+filename, img_ori)
    #         print(targretfolder+filename)
    #     else:
    #         right += 1
    #     # cv2.imshow('wrong', img_ori)
    #     # cv2.waitKey(0)
    #     # cv2.destroyAllWindows()
    #     bar.update(1)
    # print('right: %d, error: %d, acc: %s' % (right, error, right/(right+error)))